from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
#import mydata as md 




ISI = 1
DUR = 10000
BINSIZE = 50
  
L=20*umetre
dia=20*umetre
area=dia*L*math.pi
cpf=(1*ufarad*cm**-2)*area
gl=(0.1*msiemens*cm**-2)*area
glint=(0.51*msiemens*cm**-2)*area
El=-64*mV
Elint = -57.5*mV
Ek=-80*mV
Ena=55*mV
gna=(30*msiemens*cm**-2)*area
gnaint=(120*msiemens*cm**-2)*area
gk=(36*msiemens*cm**-2)*area
gkint=(100*msiemens*cm**-2)*area
gnap=(0.25*msiemens*cm**-2)*area
taumna = 0.000001*msecond
tau_m_nap = 0.01*msecond
vinit=-63
taue=5*ms
taui=5*ms
invtau=1/taue
Ee=-10*mV
Ei=-70*mV
gsynapse=50*uS*cm**-2*area

def spToList(a):
    l = list()
    
    for x in a.spiketimes:
        l+=((a[x]/20*1000/BINSIZE)/ms).tolist()
    return l

def napvtrap(x,y):
    a=x/y
    #if abs(a)<0.0001:
    #    return  1/(2+a)
    #else:
    #    return 1/(exp(a)+1)
    return 1/(exp(a)+1)
    
def napefun(x,y):
    a = x/y 
    #if abs(a)<0.0001:
    #    return 1+a
    #else:
    #    return exp(a)
    return exp(a)
            
 
def m_inf_na(V):
    return napvtrap(-(V+35),7.8)   
    
def h_inf_na(V):
    return napvtrap((V+55),7)

def tau_h_na(V):
    return msecond*(30/(napefun((V+50),15)+napefun(-(V+50),16)))

def m_inf_k(V):
    return napvtrap(-(V+28),15)

def tau_m_k(V):
    return msecond*(7/(napefun(V+40,40)+napefun(-(V+40),50)))

def m_inf_nap(V):
    return napvtrap(-V-47.1,3.1)

def h_inf_nap(V):
    return napvtrap(V+59,8)

def tau_h_nap(V):
    return msecond*(2*1200/(napefun(V+59,16)+napefun(-V-59,16)))

def i_inj(t):
    return 0*nA



# The model
eqsPF=Equations('''
dvm/dt = (-I-(gna*mna*mna*mna*hna*(vm-Ena))-(gk*mk*mk*mk*mk*(vm-Ek))-(gnap*mnap*hnap*(vm-Ena))-gl*(vm-El)+ge_current+gi_current)/cpf : mvolt
dmna/dt = (m_inf_na(vm/mvolt)-mna)/taumna : 1
dhna/dt = (h_inf_na(vm/mvolt)-hna)/tau_h_na(vm/mvolt) : 1
dmk/dt = (m_inf_k(vm/mvolt)-mk)/tau_m_k(vm/mvolt) : 1
dmnap/dt = (m_inf_nap(vm/mvolt)-mnap)/tau_m_nap : 1
dhnap/dt = (h_inf_nap(vm/mvolt)-hnap)/tau_h_nap(vm/mvolt) : 1
dge/dt = -ge*invtau : siemens
dgi/dt = -gi*invtau : siemens
dipsp/dt = (gi-ipsp)*invtau : siemens
depsp/dt = (ge-epsp)*invtau : siemens
gi_current = ipsp*(vm-Ei) : amp
ge_current = epsp*(vm-Ee) : amp

''')
#eqsPF+=alpha_conductance(input='ge',E='Ee', tau=taue,conductance_name='epsp') 
#eqsPF+=alpha_conductance(input='gi',E='Ei', tau=taui,conductance_name='ipsp') 

eqsPF+=OrnsteinUhlenbeck('I',mu=0*nA,sigma=0.002*nA,tau=10*ms)
eqsIN=Equations('''
dvm/dt = (-I-(gnaint*mna*mna*mna*hna*(vm-Ena))-(gkint*mk*mk*mk*mk*(vm-Ek))-glint*(vm-Elint)+ge_current)/cpf : mvolt
dmna/dt = (m_inf_na(vm/mvolt)-mna)/taumna : 1
dhna/dt = (h_inf_na(vm/mvolt)-hna)/tau_h_na(vm/mvolt) : 1
dmk/dt = (m_inf_k(vm/mvolt)-mk)/tau_m_k(vm/mvolt) : 1
''')
eqsIN+=alpha_conductance(input='ge',E=Ee, tau=taue,conductance_name='epsp')
eqsIN+=OrnsteinUhlenbeck('I',mu=0*nA,sigma=0.002*nA,tau=10*ms) 
print eqsIN

P=NeuronGroup(40,model=eqsPF,
threshold=EmpiricalThreshold(threshold=-40*mV),
implicit=True)

PE=P.subgroup(20)
PF=P.subgroup(20)

FREQ=30*Hz
spiketimes = [(0,5*ms),(1,5*ms),(2,5*ms),(3,5*ms),(4,5*ms),(5,5*ms),(6,5*ms),(7,5*ms),(8,5*ms),(9,5*ms),(10,5*ms),(11,5*ms),(12,5*ms),(13,5*ms),(14,5*ms),(15,5*ms),(16,5*ms),(17,5*ms),(18,5*ms),(19,5*ms),]

print area
print gsynapse
 
SG = SpikeGeneratorGroup(20,spiketimes,period=1/FREQ)
CINPUT1=Connection(SG,PE,'ge',weight=gsynapse*0.08,sparseness=1,delay=lambda i,j:15*ms+randn(len(PE))*4*ms);
CINPUT2=Connection(SG,PF,'ge',weight=gsynapse*0.078,sparseness=1,delay=lambda i,j:15*ms+randn(len(PE))*4*ms);

IN=NeuronGroup(40,model=eqsIN,
threshold=EmpiricalThreshold(threshold=-20*mV),
implicit=True)

INE=IN.subgroup(20)
INF=IN.subgroup(20)


#connections between Pattern generation neurons to inhibitory interneurons
CinE=Connection(PE,INF,'ge',weight=gsynapse*0.45,sparseness=1,delay=lambda i,j:7*ms+randn(len(PE))*2*ms);
CinF=Connection(PF,INE,'ge',weight=gsynapse*0.45,sparseness=1,delay=lambda i,j:7*ms+randn(len(PE))*2*ms);

#self potentiation of pattern generation neurons
Cs1=Connection(PE,PE,'ge',weight=gsynapse*0.0125,sparseness=1,delay=lambda i,j:2*ms+randn(len(PE))*0.3*ms);
Cs2=Connection(PF,PF,'ge',weight=gsynapse*0.0125,sparseness=1,delay=lambda i,j:2*ms+randn(len(PF))*0.3*ms);

#reciprocal potentiation of pattern generation neurons
Cr1=Connection(PE,PF,'ge',weight=gsynapse*0.0125,sparseness=1,delay=lambda i,j:14*ms+randn(len(PE))*0.6*ms);
Cr2=Connection(PF,PE,'ge',weight=gsynapse*0.0125,sparseness=1,delay=lambda i,j:14*ms+randn(len(PF))*0.6*ms);

#inhibition of pattern generation neurons
CinE2=Connection(INE,PE,'gi',weight=gsynapse*0.115,sparseness=1,delay=lambda i,j:7*ms+randn(len(PE))*2*ms);
CinF2=Connection(INF,PF,'gi',weight=gsynapse*0.115,sparseness=1,delay=lambda i,j:7*ms+randn(len(PE))*2*ms);




# Initialization

vv =vinit
P.vm=vv*mV 
P.mna = m_inf_na(vv)
P.hna = h_inf_na(vv)
P.mk = m_inf_k(vv)

P.mnap = m_inf_nap(vv)
P.hnap = h_inf_nap(vv)
#P.El = -64*mV * randn(len(P))*0.64*mV


vin =-64
IN.vm=vin*mV
IN.mna = m_inf_na(vin)
IN.hna = h_inf_na(vin)
IN.mk = m_inf_k(vin);
#IN.El = -57.5*mV * randn(len(P))*2.875*mV

# Record a few trace
tracePF=StateMonitor(PF,'vm',record=[0])
tracePFgi=StateMonitor(PF,'gi_current',record=[0])
tracePFge=StateMonitor(PF,'ge_current',record=[0])
tracePE=StateMonitor(PE,'vm',record=[0])
traceINE=StateMonitor(INE,'vm',record=[0])
traceINF=StateMonitor(INF,'vm',record=[0])

spikesPF = SpikeMonitor(PF)
spikesPE = SpikeMonitor(PE)
spikesINE= SpikeMonitor(INE)
spikesINF= SpikeMonitor(INF)

run(DUR*msecond)
figure(1)
subplot(411)
plot(tracePF.times/ms,tracePF[0]/mV)

subplot(412)
plot(tracePE.times/ms,tracePE[0]/mV)
#plot(tracePFgi.times/ms,tracePFgi[0]/nA)
#plot(tracePFge.times/ms,tracePFge[0]/nA)
subplot(413)
plot(traceINE.times/ms,traceINE[0]/mV)
subplot(414)
plot(traceINF.times/ms,traceINF[0]/mV)



figure(2)
subplot(411)
#raster_plot(spikesPE)
x = spToList(spikesPF)
hist(x,floor(DUR/BINSIZE),range=[0,DUR],histtype='step')

subplot(412)
x = spToList(spikesPE)
hist(x,floor(DUR/BINSIZE),range=[0,DUR],histtype='step')
subplot(413)
x = spToList(spikesINE)
hist(x,floor(DUR/BINSIZE),range=[0,DUR],histtype='step')
subplot(414)
x = spToList(spikesINF)
hist(x,floor(DUR/BINSIZE),range=[0,DUR],histtype='step')
show()